import time

import numpy as np
import pandas as pd
from numba import njit, vectorize


# 定义一个加速函数
# 0.225445 秒
@njit
def custom_function(x):
    return x**2 + np.sin(x)


# 1.123084 秒
def custom_function2(x):
    return x**2 + np.sin(x)


# 并行加速
# vectorize 专门用来处理逐元素操作，且可以支持多种数据类型。
# guvectorize对固定形状的块进行操作，适合更复杂的广播逻辑。
# 0.006977 秒
@vectorize(['float32(float32)', 'float64(float64)'], target='parallel')
def custom_function3(x):
    return x**2 + np.sin(x)


# 创建一个 DataFrame
df = pd.DataFrame({'A': np.random.rand(10**6)})

custom_function(np.random.rand())

t1 = time.time()
# 使用 Numba 加速 apply
df['B'] = df['A'].apply(custom_function)
t2 = time.time()
print(f'时长: {t2-t1:.6f} 秒')

t1 = time.time()
df['C'] = df['A'].apply(custom_function2)
t2 = time.time()
print(f'时长: {t2-t1:.6f} 秒')

custom_function3(np.random.rand())
t1 = time.time()
# 使用 Numba 加速 apply
df['D'] = df['A'].apply(custom_function3)
t2 = time.time()
print(f'时长: {t2-t1:.6f} 秒')
